package com.campus.counseling.service.impl;

import com.campus.counseling.config.LSTMConfig;
import com.campus.counseling.config.ModelTrainingConfig;
import com.campus.counseling.service.LSTMModelTrainer;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.LSTM;
import org.deeplearning4j.nn.conf.layers.RnnOutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;
import org.springframework.stereotype.Service;
import org.deeplearning4j.util.ModelSerializer;

import javax.annotation.PostConstruct;
import java.io.File;
import java.util.Arrays;
import java.nio.file.Files;
import java.nio.file.StandardCopyOption;

@Slf4j
@Service
@RequiredArgsConstructor
public class LSTMModelTrainerImpl implements LSTMModelTrainer {

    private final LSTMConfig lstmConfig;
    private final ModelTrainingConfig trainingConfig;
    private MultiLayerNetwork model;
    private double bestScore = Double.MAX_VALUE;

    @PostConstruct
    public void init() {
        try {
            model = loadModel();
        } catch (Exception e) {
            log.warn("模型加载失败，将在训练时初始化新模型");
        }
    }

    @Override
    public void initModel() {
        try {
            log.info("开始初始化LSTM模型...");
            
            // 创建一个NeuralNetConfiguration.Builder对象，用于构建神经网络配置
            NeuralNetConfiguration.Builder builder = new NeuralNetConfiguration.Builder()
                .seed(123) // 设置随机种子
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT) // 设置优化算法
                .updater(new Adam(0.01)) // 设置更新器
                .weightInit(WeightInit.XAVIER) // 设置权重初始化方法
                .l2(1e-5); // 设置L2正则化参数

            // 创建一个NeuralNetConfiguration.ListBuilder对象，用于构建神经网络层列表
            NeuralNetConfiguration.ListBuilder listBuilder = builder.list();

            // 简化模型结构，使用单层LSTM
            listBuilder.layer(0, new LSTM.Builder()
                .nIn(lstmConfig.getInputSize()) // 设置输入层大小
                .nOut(lstmConfig.getHiddenSize()) // 设置隐藏层大小
                .activation(Activation.TANH) // 设置激活函数
                .dropOut(0.5) // 设置dropout率
                .build());

            // 输出层
            listBuilder.layer(1, new RnnOutputLayer.Builder()
                .nIn(lstmConfig.getHiddenSize()) // 设置输入层大小
                .nOut(lstmConfig.getOutputSize()) // 设置输出层大小
                .activation(Activation.SIGMOID) // 设置激活函数
                .lossFunction(LossFunctions.LossFunction.XENT) // 设置损失函数
                .build());

            // 构建多层神经网络配置
            MultiLayerConfiguration config = listBuilder.build();
            // 创建多层神经网络模型
            model = new MultiLayerNetwork(config);
            // 初始化模型
            model.init();
            
            log.info("LSTM模型初始化完成，结构:\n{}", model.summary());

        } catch (Exception e) {
            log.error("LSTM模型初始化失败:", e);
            throw new RuntimeException("LSTM模型初始化失败", e);
        }
    }

    @Override
    public void train(DataSet trainingData) {
        if (model == null) {
            initModel();  // 确保模型已初始化
        }

        try {
            log.info("开始训练LSTM模型...");
            
            // 获取原始数据
            INDArray features = trainingData.getFeatures();
            INDArray labels = trainingData.getLabels();
            
            // 调整特征形状为 [batchSize, inputSize, timeSteps]
            features = features.permute(0, 2, 1);
            
            // 调整标签形状为 [batchSize, outputSize, timeSteps]
            labels = labels.reshape(labels.size(0), 1, 1);
            
            log.info("训练数据形状: features={}, labels={}",
                Arrays.toString(features.shape()),
                Arrays.toString(labels.shape()));

            // 添加训练监听器
            model.setListeners(new ScoreIterationListener(1));
            
            // 训练多个epoch
            int numEpochs = trainingConfig.getNumEpochs();
            for (int epoch = 0; epoch < numEpochs; epoch++) {
                // 训练一个epoch
                model.fit(new DataSet(features, labels));
                
                // 计算当前loss
                double score = model.score();
                log.info("Epoch {}/{}, loss: {}", epoch + 1, numEpochs, score);
                
                // 如果是最佳模型则保存
                if (score < bestScore) {
                    bestScore = score;
                    String bestModelPath = lstmConfig.getModelPath().replace(".bin", "_best.bin");
                    ModelSerializer.writeModel(model, bestModelPath, true);
                    log.info("保存最佳模型，得分: {}", score);
                }
            }
            
            // 保存最终模型
            saveModel();
            log.info("LSTM模型训练完成，最终loss: {}", model.score());
            
        } catch (Exception e) {
            log.error("LSTM模型训练失败:", e);
            throw new RuntimeException("LSTM模型训练失败", e);
        }
    }

    @Override
    public INDArray predict(INDArray features) {
        // 判断模型是否初始化
        if (model == null) {
            log.warn("模型未初始化，无法进行预测");
            // 返回全零数组
            return Nd4j.zeros(features.size(0), 1);
        }

        try {
            // 调整特征形状为 [batchSize, inputSize, timeSteps]
            features = features.permute(0, 2, 1);
            
            // 进行预测
            INDArray output = model.output(features);
            // 将输出从 [batchSize, outputSize, timeSteps] 转换为 [batchSize, outputSize]
            output = output.tensorAlongDimension(0, 0, 1);
            
            log.debug("预测结果形状: {}", Arrays.toString(output.shape()));
            return output;
        } catch (Exception e) {
            log.error("预测失败", e);
            // 返回全零数组
            return Nd4j.zeros(features.size(0), 1);
        }
    }

    @Override
    public void saveModel() {
        try {
            // 先备份旧模型
            File oldModel = new File(lstmConfig.getModelPath());
            if (oldModel.exists()) {
                String backupPath = lstmConfig.getModelPath() + ".backup";
                Files.copy(oldModel.toPath(), new File(backupPath).toPath(), StandardCopyOption.REPLACE_EXISTING);
                log.info("已备份旧模型到: {}", backupPath);
            }
            
            // 保存新模型
            String modelPath = lstmConfig.getModelPath();
            log.info("开始保存模型到: {}", modelPath);
            
            File modelFile = new File(modelPath);
            modelFile.getParentFile().mkdirs();
            
            ModelSerializer.writeModel(model, modelPath, true);
            log.info("模型保存成功");
            
            // 保存最佳模型
            double score = model.score();
            String bestModelPath = modelPath.replace(".bin", "_best.bin");
            if (!new File(bestModelPath).exists() || score < bestScore) {
                ModelSerializer.writeModel(model, bestModelPath, true);
                bestScore = score;
                log.info("保存最佳模型，得分: {}", score);
            }
            
        } catch (Exception e) {
            log.error("保存模型失败: ", e);
            throw new RuntimeException("保存模型失败: " + e.getMessage());
        }
    }

    @Override
    public MultiLayerNetwork loadModel() {
        try {
            // 获取模型路径
            String modelPath = lstmConfig.getModelPath();
            log.info("开始加载模型: {}", modelPath);
            
            // 创建模型文件对象
            File modelFile = new File(modelPath);
            // 判断模型文件是否存在
            if (!modelFile.exists()) {
                // 如果不存在，抛出异常
                throw new RuntimeException("模型文件不存在: " + modelPath);
            }
            
            // 使用ModelSerializer.restoreMultiLayerNetwork方法加载模型
            MultiLayerNetwork loadedModel = ModelSerializer.restoreMultiLayerNetwork(modelFile);
            log.info("模型加载成功");
            // 将加载的模型赋值给成员变量
            this.model = loadedModel;
            // 返回加载的模型
            return loadedModel;
            
        } catch (Exception e) {
            // 如果加载模型失败，记录错误日志，并抛出异常
            log.error("加载模型失败: ", e);
            throw new RuntimeException("加载模型失败: " + e.getMessage());
        }
    }
} 